
import numpy as np


def create_meshgrid(h, w, step=7, patch_size=14, return_hw=False):
        start_coord = patch_size//2
        x = torch.arange(start_coord, w, step=step, device=device).float()
        y = torch.arange(start_coord, h, step=step, device=device).float()
        yy, xx = torch.meshgrid(y, x)
        xx = xx.reshape(-1)
        yy = yy.reshape(-1)
        grid = torch.stack([xx, yy], dim=-1)
        if return_hw:
            return grid, len(y), len(x)
        return grid

def test_create_meshgrid():
    h=476; w=854
    grid, ly, lx = create_meshgrid(h=h, w=w, step=7, patch_size=14, return_hw=True)

    print(h*w/7/7)
    print(grid.shape)
    print(ly, lx)


    print((np.array([1270,720]) - (14 // 2)) / 7)


    print(67*121)
    # plot the grid
    import matplotlib.pyplot as plt
    plt.scatter(grid[:,0], grid[:,1])
    plt.show()

import torch
from torchvision.ops import batched_nms #, nms
def get_bb_sim_indices(affs_batched, coords, box_size=10, iou_thresh=0.3, topk=3, device=device):
    """  affs_batched: B x N """
    topk = torch.topk(affs_batched, k=topk, sorted=False, dim=1)
    filt_idx = topk.indices # B x topk
    affs_filt = topk.values # B x topk

    if affs_filt.shape[0] == 0:
        return None, None, None
    
    filt_coords = coords[filt_idx]
    xmin = filt_coords[:, :, 0] - box_size # B x topk
    xmax = filt_coords[:, :, 0] + box_size # B x topk
    ymin = filt_coords[:, :, 1] - box_size # B x topk
    ymax = filt_coords[:, :, 1] + box_size # B x topk
    # concat to get boxes shaped B x topk x 4
    boxes = torch.cat([xmin[:, :, None], ymin[:, :, None], xmax[:, :, None], ymax[:, :, None]], dim=-1) # B x topk x 4
    # get idxs shaped (B x topk) representing the batch index
    idxs = torch.arange(filt_idx.shape[0], device=device)[:, None].repeat(1, filt_idx.shape[1]).reshape(-1) # (B x topk)
    peak_indices = batched_nms(boxes.float().reshape(-1, 4), affs_filt.reshape(-1), idxs, iou_thresh)
    # convert peak_indices to the original indices to the  not flat indices
    peak_indices_original = torch.stack([peak_indices // filt_idx.shape[1], peak_indices % filt_idx.shape[1]], dim=-1)
    # retrieve the first two elements of the peak_indices_original for the first axis
    filt_idx_mask = torch.zeros_like(filt_idx, device=device) # B x topk
    filt_idx_mask[peak_indices_original[:, 0], peak_indices_original[:, 1]] = 1
    peak_aff_batched = affs_filt * filt_idx_mask # B x topk
    # retrieve the highest and second highest affinities for each batch
    top2 = torch.topk(peak_aff_batched, k=2, dim=1)
    top2_values, top2_indices = top2.values, top2.indices # B x 2, B x 2
    highest_affs, highest_affs_idx = top2_values[:, 0], top2_indices[:, 0] # B, B
    second_highest_affs, second_highest_affs_idx = top2_values[:, 1], top2_indices[:, 1] # B, B
    r = second_highest_affs / highest_affs 
    # 筛选出的数据将具有形状 (B, 2)
    top2_filt_idx = torch.gather(filt_idx, 1, top2_indices)
    
    return top2_filt_idx, top2_values, r


def test_get_bb_sim_indices():
    # # 定义一个相似度矩阵，模拟存在局部干扰的情况
    # affinity = torch.tensor([
    #     [0.1, 0.8, 0.2, 0.3, 0.4],
    #     [0.3, 0.4, 0.93, 0.91, 0.81],
    #     [0.7, 0.95, 0.5, 0.1, 0.92]
    # ])

    # # 假设的目标点坐标
    # coords = torch.tensor([
    #     [7, 7],
    #     [14, 7],
    #     [21, 7],
    #     [21, 14],
    #     [21, 21]
    # ])
    # 构造相似度矩阵 affinity
    # 构造相似度矩阵 affinity
    affinity = torch.tensor([
        [0.9, 0.89, 0.85, 0.8, 0.88],  # 行0：最高分在列0（坐标10,10）
        [0.95, 0.90, 0.93, 0.92, 0.91], # 行1：最高分在列0（坐标10,10）
        [0.92, 0.91, 0.90, 0.93, 0.88], # 行2：最高分在列0（坐标10,10）
        [0.8, 0.805, 0.6, 0.81, 0.4]     # 行3：最高分在列0（坐标10,10）
    ])

    coords = torch.tensor([
        [10, 11],  # 列0
        [11, 11],  # 列1
        [20, 20],  # 列1
        [10, 12] ,  # 列1
        [40, 40] 
    ])

    # 假设的批量相似度数据
    affs_batched = affinity

    # 调用 get_bb_sim_indices 函数
    top2_filt_idx, top2_values, r = get_bb_sim_indices(affs_batched, coords)
    highest_indices_method1 = top2_filt_idx[:, 0]
    print("方法一（get_bb_sim_indices）选出的最高相似度坐标索引：", highest_indices_method1)
    affinity_source_max = torch.argmax(affinity, dim=1)  # n
    print("affinity_source_max:", affinity_source_max) # 121
    highest_affs_idx =  top2_filt_idx[:, 1] 
    print("方法一（get_bb_sim_indices）选出的次高相似度坐标索引：", highest_affs_idx)

    # # 第二段代码的操作
    # feature_range = torch.arange(affinity.shape[0])
    # affinity_source_max = torch.argmax(affinity, dim=1)  # n
    # print("affinity_source_max:", affinity_source_max) # 121
    # affinity_target_max = torch.argmax(affinity, dim=0)  # m
    # print("affinity_target_max:", affinity_target_max) # 221
    # affinity_target_max = torch.argmax(affinity, dim=0)  # m
    # source_bb_indices = feature_range == affinity_target_max[affinity_source_max] # 12
    # print("source_bb_indices:", source_bb_indices)
    # target_bb_indices = affinity_source_max[source_bb_indices]

    # print("方法二（直接 argmax）选出的最高相似度坐标索引：", target_bb_indices)
    # affinities = affinity[feature_range[source_bb_indices], target_bb_indices]
    # print("方法二（直接 argmax）选出的最高相似度值：", affinities)

